import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from PyEMD import EMD, EEMD, Visualisation # 导入PyEMD库
from scipy.signal import hilbert # 用于希尔伯特变换
import numpy as np

# --- 1. 参数配置 ---
# 请确保这个路径与你的Excel文件实际存放的路径完全一致
file_path = r"D:\Mine in WPS\流致振动\1.xlsx"

# 定义要分析的测点和时间窗口
section_to_analyze = 'A' # 可以改为 'B' 或 'C'
time_interval = (15, 24) # 分析的时间段，例如 Ramp-up 阶段

# --- 2. 加载和预处理数据 ---
try:
    # 读取Excel文件，假设没有表头，第一列为时间，2/3/4列为数据
    df = pd.read_excel(file_path, skiprows=1, names=['Time', 'A', 'B', 'C'])
    # 将 'Time' 列转换为数值类型，以便进行比较
    df['Time'] = pd.to_numeric(df['Time'])
    print("文件加载成功，数据前5行：")
    print(df.head())
except FileNotFoundError:
    print(f"错误：找不到文件 '{file_path}'。请检查路径是否正确。")
    exit()

# 为了更好地显示中文和负号
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# --- 3. 数据切片 ---
start_time, end_time = time_interval
df_slice = df[(df['Time'] >= start_time) & (df['Time'] <= end_time)]

if df_slice.empty:
    print(f"警告：在时间段 ({start_time}s - {end_time}s) 没有找到数据。")
    exit()

# 获取位移数据和时间数据
displacement = df_slice[section_to_analyze]
time = df_slice['Time']
N = len(time)

print(f"分析 {section_to_analyze} 截面在 {start_time}s 到 {end_time}s 的数据，共 {N} 个点。")

# --- 4. HHT分析 ---

# 4.1 经验模态分解 (EMD)
print("开始进行EMD分解...")
emd = EMD() # 创建EMD实例
imfs = emd(displacement.to_numpy()) # 对位移信号进行EMD分解
# imfs 是一个二维数组，每一行是一个IMF，最后一行是残余趋势
num_imfs = imfs.shape[0]
print(f"EMD分解完成，共得到 {num_imfs} 个 IMF 和 1 个 Residue。")

# 4.2 希尔伯特谱分析 (HSA)
print("开始进行希尔伯特变换计算瞬时频率和振幅...")
# 准备存储瞬时频率和振幅的列表
instantaneous_frequencies = []
instantaneous_amplitudes = []

# 对每个IMF（除了最后一个趋势项）进行希尔伯特变换
for i in range(num_imfs - 1): # 最后一个通常是趋势，不分析
    imf = imfs[i]
    # 希尔伯特变换
    analytic_signal = hilbert(imf)
    # 瞬时振幅
    amplitude_envelope = np.abs(analytic_signal)
    # 瞬时相位
    instantaneous_phase = np.unwrap(np.angle(analytic_signal))
    # 瞬时频率 (对时间求导，注意单位是 rad/s，转换为 Hz 需要除以 2*pi)
    # np.gradient 对于非均匀采样可能不准确，这里假设近似均匀
    inst_freq = (np.gradient(instantaneous_phase) / (2.0*np.pi)) * (N / (time.iloc[-1] - time.iloc[0]))
    
    instantaneous_frequencies.append(inst_freq)
    instantaneous_amplitudes.append(amplitude_envelope)

print("希尔伯特变换完成。")

# --- 5. 结果可视化 ---

# 5.1 绘制原始信号和IMFs
fig_imfs, axes_imfs = plt.subplots(num_imfs + 1, 1, figsize=(12, 2 * (num_imfs + 1)), sharex=True)
fig_imfs.suptitle(f'{section_to_analyze} 截面 EMD 分解结果', fontsize=16)

# 绘制原始信号
axes_imfs[0].plot(time, displacement, 'b-', linewidth=1)
axes_imfs[0].set_ylabel('原始信号')
axes_imfs[0].grid(True)

# 绘制IMFs
for i in range(num_imfs):
    axes_imfs[i+1].plot(time, imfs[i], 'r-', linewidth=1)
    if i < num_imfs - 1:
        axes_imfs[i+1].set_ylabel(f'IMF-{i+1}')
    else:
        axes_imfs[i+1].set_ylabel('Residue')
    axes_imfs[i+1].grid(True)

axes_imfs[-1].set_xlabel('时间 (s)')
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
# plt.show() # 暂时注释掉show
fig_imfs.savefig(f'{section_to_analyze}_IMFs.png') # 保存图像
print(f"IMFs图已保存为 {section_to_analyze}_IMFs.png")


# 5.2 绘制希尔伯特谱 (Hilbert Spectrum)
print("开始绘制希尔伯特谱...")
# 创建网格用于绘图
t, f = np.meshgrid(time, np.linspace(0, 15, 300)) # 假设我们关注0-15Hz
hht_magnitude = np.zeros_like(t)

# 将每个IMF的时频信息叠加到Hilbert谱矩阵中
# 这里采用一种简化的方法：对于每个时间点，找到其瞬时频率最接近的网格f点，然后累加振幅
for i in range(len(instantaneous_frequencies)):
    for j in range(N):
        freq_val = instantaneous_frequencies[i][j]
        amp_val = instantaneous_amplitudes[i][j]
        # 找到最接近的频率网格点
        if 0 <= freq_val <= 15: # 限制在0-15Hz范围内
            freq_idx = np.argmin(np.abs(f[:,j] - freq_val))
            hht_magnitude[freq_idx, j] += amp_val # 累加振幅

# 绘制Hilbert谱
fig_hht, ax_hht = plt.subplots(figsize=(12, 6))
# 使用pcolormesh绘制，注意X和Y的维度
cax = ax_hht.pcolormesh(t, f, hht_magnitude, shading='auto', cmap='jet')
ax_hht.set_xlabel('时间 (s)')
ax_hht.set_ylabel('频率 (Hz)')
ax_hht.set_title(f'{section_to_analyze} 截面 Hilbert Spectrum')
ax_hht.set_ylim(0, 15) # 限制频率显示范围
fig_hht.colorbar(cax, ax=ax_hht, label='瞬时振幅')
plt.tight_layout()
# plt.show() # 暂时注释掉show
fig_hht.savefig(f'{section_to_analyze}_HilbertSpectrum.png') # 保存图像
print(f"Hilbert谱图已保存为 {section_to_analyze}_HilbertSpectrum.png")

# 5.3 绘制边际谱 (Marginal Spectrum)
print("开始绘制边际谱...")
marginal_spectrum = np.sum(hht_magnitude, axis=1) # 对时间轴求和
freq_axis_for_marginal = f[:,0] # 频率轴

fig_marginal, ax_marginal = plt.subplots(figsize=(8, 5))
ax_marginal.plot(freq_axis_for_marginal, marginal_spectrum, 'g-', linewidth=1.5)
ax_marginal.set_xlabel('频率 (Hz)')
ax_marginal.set_ylabel('能量')
ax_marginal.set_title(f'{section_to_analyze} 截面 Marginal Spectrum')
ax_marginal.grid(True)
ax_marginal.set_xlim(0, 15)
plt.tight_layout()
# plt.show() # 暂时注释掉show
fig_marginal.savefig(f'{section_to_analyze}_MarginalSpectrum.png') # 保存图像
print(f"边际谱图已保存为 {section_to_analyze}_MarginalSpectrum.png")

print("HHT分析及可视化完成。")